import numpy as np

from paddle.fluid.io import Dataset
from helm.datasets.mnist import MNIST
from helm.datasets.cifar import CIFAR10, CIFAR100


class Fullset(Dataset):

    def __init__(self, dataset, transform):
        super().__init__()
        self.dataset = dataset
        self.transform = transform

    def __getitem__(self, idx):
        input, target = self.dataset[idx]
        return self.transform(input, target)

    def __len__(self):
        return len(self.dataset)

    def __repr__(self):
        return "Fullset(%s)" % self.dataset


class Subset(Dataset):
    """
    Subset of a dataset at specified indices.

    Arguments:
        dataset (Dataset): The whole Dataset
        indices (sequence): Indices in the whole set selected for subset
    """

    def __init__(self, dataset, indices, transform=None):
        super().__init__()
        self.dataset = dataset
        self.indices = indices
        self.transform = transform

    def __getitem__(self, idx):
        img, target = self.dataset[self.indices[idx]]

        if self.transform is not None:
            img = self.transform(img)

        return img, target

    def get_image(self, idx):
        return self.dataset.get_image(self.indices[idx])

    def get_target(self, idx):
        return self.dataset.get_target(self.indices[idx])

    def __len__(self):
        return len(self.indices)

    def __repr__(self):
        fmt_str = 'Subset of ' + self.dataset.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        return fmt_str


def train_test_split(dataset, test_ratio, random=False, transform=None, test_transform=None):
    num_examples = len(dataset)
    num_test_examples = int(num_examples * test_ratio)
    num_train_examples = num_examples - num_test_examples
    if random:
        indices = np.random.permutation(num_examples)
        train_indices = indices[:num_train_examples]
        test_indices = indices[num_train_examples:]
    else:
        train_indices = np.arange(num_train_examples)
        test_indices = np.arange(num_train_examples, num_examples)
    train_set = Subset(
        dataset, train_indices, transform)
    test_set = Subset(
        dataset, test_indices, test_transform)
    return train_set, test_set


class CombineDataset(Dataset):

    def __init__(self, *datasets):
        super().__init__()
        self.datasets = datasets
        assert len(set(len(ds) for ds in datasets)) == 1, "All datasets must be of the same length"

    def __len__(self):
        return len(self.datasets[0])

    def __getitem__(self, idx):
        return tuple(x for ds in self.datasets for x in ds[idx])